# 3rd
import pandas as pd

# std
import os
import datetime
import pickle

import getpass

if os.access('/remote/iosg/coin/data/order_latency', os.W_OK):
  DATA_ROOT = '/remote/iosg/coin/data/order_latency'
else:
  DATA_ROOT = f'/remote/iosg/home-2/{getpass.getuser()}/order_latency'

ALL_COLS = [
    'order_machine',
    'feed_machine',
    'market_type',
    'exchange',
    'symbol',
    'timestamp',
    'price',
    'qty',
    'latency'
]


class LatencyModel(object):
  def __init__(self, alpha=0.5):
    """[summary]

    Arguments:
        object {[type]} -- [description]

    Keyword Arguments:
        alpha {[type]} -- a regularization coefficient for Ridge regression
    """
    self.model = None
    self.cols = None
    self.dummy_cols = None

  def predict(self, record, nondummy_dict=None):
    """

    predict({'exchange': 'Bitmex', 'market_type': 'Futures'})

    pandas.Series
      exchange_Bitmex exchange_Okex ..
      1               0             ..

    Arguments:
        record {pandas.Series or pandas.DataFrame or dict} -- a single record or DataFrame
    """

    if self.model is None:
      print('[error] model is not prepared.')
      return

    is_single_record = False

    if isinstance(record, dict):
      keys = record.keys()
      for key in keys:
        if key not in self.cols:
          print(f'[warning] the model has no column {key}. ignoring this column.')
          print(f'available columns: {self.cols}')

      for key in self.cols:
        if key not in keys and (nondummy_dict is None or key not in nondummy_dict):
          print(f'[warning] the model requires a column {key}. The prediction will be inaccurate.')

      record = pd.Series(record)

    if isinstance(record, pd.Series):
      is_single_record = True
      df = record.to_frame().T

    df = pd.get_dummies(df)
    df = df.reindex(columns=self.dummy_cols, fill_value=0)

    if nondummy_dict is not None:
      for key, value in nondummy_dict.items():
        df[key] = value

    ret = self.model.predict(df)
    if is_single_record:
      ret = ret[0]

    return ret

  def read_one_day(self, d, drop_cols):
    dstr = datetime.datetime.strftime(d, '%Y%m%d')
    path = f'{DATA_ROOT}/dataframes/{dstr}.h5'
    print('reading file:', path)
    df = pd.read_hdf(path, 'df')

    if len(df) > 0:
      df.drop(drop_cols, axis=1, inplace=True)

    return df

  def train_for_dates(self, dates, cols, alpha=0.5, ycol='latency'):
    """[summary]

    Arguments:
        dates {str} -- a comma-sepearted list of date ranges or a single day.
                       e.g. 20190502,20190814-20190816
        cols {str} -- a comma-separated list of columns
    """

    assert isinstance(dates, str)
    assert isinstance(cols, str)

    # procoess cols
    cols = cols.split(',')
    for col in cols:
      if col not in ALL_COLS and ycol == 'latency':
        print(f'column {col} does not exist')
        import sys
        sys.exit(1)
    if ycol not in cols and ycol == 'latency':
      cols.append(ycol)
    drop_cols = [col for col in ALL_COLS if col not in cols]

    print(' using cols:', cols)
    print('unused cols:', drop_cols)

    ##########

    # process dates
    def parse_dates(d):  # d="YYYYmmdd" or d="YYYYmmdd-YYYYmmdd"
      if len(d) == len('20001111'):
        d = datetime.datetime.strptime(d, '%Y%m%d')
        return (d, d)

      d1 = datetime.datetime.strptime(d[:8], '%Y%m%d')
      d2 = datetime.datetime.strptime(d[9:], '%Y%m%d')
      return (d1, d2)

    date_pairs = [parse_dates(d) for d in dates.split(',')]
    ##########

    # get dataframes
    frames = []
    for (s, e) in date_pairs:
      day = s
      while True:
        print(f'reading data for {day} ..')
        frame = self.read_one_day(day, drop_cols)
        if frame is not None and len(frame) > 0:  # broken data?
          frames.append(frame)

        if day == e:
          break
        day += datetime.timedelta(days=1)
    ##########

    # cleanup
    print('aggregating ..')
    df = pd.concat(frames)
    del frames

    if ycol == 'latency':
      print('filtering outliers ..')
      df = df[df.latency < 1e9]
      ##########

    self.train_for_df(df, alpha, cols)

  def train_for_df(self, df, alpha=0.5, cols=None):
    from sklearn.linear_model import Ridge
    """[summary]

    Arguments:
        df {pandas.DataFrame} -- a dataframe with latency column
    """

    assert isinstance(df, pd.DataFrame)

    y = df['latency']
    x = df.drop('latency', axis=1)
    self.cols = x.columns

    x = pd.get_dummies(x)
    self.dummy_cols = x.columns

    self.model = Ridge(alpha=alpha)
    self.model.fit(x, y)

  def load(self, path):
    with open(path, 'rb') as f:
      self.model = pickle.load(f)
      self.cols = pickle.load(f)
      self.dummy_cols = pickle.load(f)

  def save(self, path):
    with open(path, 'wb') as f:
      pickle.dump(self.model, f)  # Ridge params
      pickle.dump(self.cols, f)
      pickle.dump(self.dummy_cols, f)  # one-hot encoding info
